package org.onion.pattern.kit.common.lock;


import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.locks.ReentrantLock;

/**
 * 分段锁，系统提供一定数量的原始锁，根据传入对象的哈希值获取对应的锁并加锁 <br>
 * 注意：要锁的对象的哈希值如果发生改变，有可能导致锁无法成功释放!!!
 */
public class SegmentLock<T> {
    private int segments = 16;//默认分段数量
    private final Map<Integer, ReentrantLock> lockMap;

    public SegmentLock() {
        this(false);
    }

    public SegmentLock(boolean fair) {
        this(fair, 16);
    }

    /**
     * @param counts 范围在 3-30 之间
     * @param fair
     */
    public SegmentLock(boolean fair, int counts) {
        segments = counts;
        lockMap = new ConcurrentHashMap<>(segments);
        for (int i = 0; i < segments; i++) {
            lockMap.put(i, new ReentrantLock(fair));
        }
    }

    /**
     * 加锁
     *
     * @param key
     */
    public void lock(T key) {
        ReentrantLock lock = lockMap.get((key.hashCode() >>> 1) % segments);
        lock.lock();
    }

    public void lockInterruptibly(T key) throws InterruptedException {
        ReentrantLock lock = lockMap.get((key.hashCode() >>> 1) % segments);
        lock.lockInterruptibly();
    }

    public void tryLock(T key) {
        ReentrantLock lock = lockMap.get((key.hashCode() >>> 1) % segments);
        lock.tryLock();
    }

    public void tryLock(T key, long timeout, TimeUnit unit) throws InterruptedException {
        ReentrantLock lock = lockMap.get((key.hashCode() >>> 1) % segments);
        lock.tryLock(timeout, unit);
    }

    /**
     * 释放锁
     *
     * @param key
     */
    public void unlock(T key) {
        ReentrantLock lock = lockMap.get((key.hashCode() >>> 1) % segments);
        lock.unlock();
    }

    /**
     * 获取锁对象
     *
     * @param key
     * @return
     */
    public ReentrantLock getLock(T key) {
        return lockMap.get((key.hashCode() >>> 1) % segments);
    }


    public int getLockSize() {
        return lockMap.size();
    }
}
